'''
Created on Aug 17, 2012

@author: issa
'''
import numpy as np
import cv2
from common import anorm , clock
from functools import partial
import time
import video
SKIPFRAME = 4
FLANN_INDEX_KDTREE = 1  # bug: flann enums are missing

flann_params = dict(algorithm = FLANN_INDEX_KDTREE,
                    trees = 4)


class Tracker (object):

    def __init__(self , tkp , tdesc , skp , sdesc , threshold):
        self.tdesc = tdesc
        self.sdesc = sdesc
        self.tkp = tkp
        self.skp = skp 
        self.threshold = threshold
        
    def match_bruteforce(self):
        res = []
        for i in xrange(len(self.tdesc)):
            dist = anorm( self.sdesc - self.tdesc[i] )
            n1, n2 = dist.argsort()[:2]
            r = dist[n1] / dist[n2]
            if r < self.threshold:
                res.append((i, n1))
                return np.array(res)
    
    def match_flann(self):
        flann = cv2.flann_Index(self.sdesc, flann_params)
        idx2, dist = flann.knnSearch(self.tdesc, 2, params = {}) # bug: need to provide empty dict
        mask = dist[:,0] / dist[:,1] < self.threshold
        idx1 = np.arange(len(self.tdesc))
        pairs = np.int32( zip(idx1, idx2[:,0]) )
        return pairs[mask]
    



if __name__ =="__main__":
    fn1 = 'issa.png' 
    im1 = cv2.imread(fn1, 0)
    itr = 0
    cap = video.create_capture(1)
    surf_det = cv2.FeatureDetector_create("SURF")
    surf_ext = cv2.DescriptorExtractor_create("SURF")
    kp1 = surf_det.detect(im1)
    kp1, desc1 = surf_ext.compute(im1, kp1)
    while True:
        tt = clock()
        ret , img2 = cap.read()

        if itr % SKIPFRAME ==0 :
            img3 = cv2.blur(img2, (7,7))
            img3 = cv2.cvtColor(img3, cv2.COLOR_RGB2GRAY)
            kp2 = surf_det.detect(img3)
            kp2, desc2 = surf_ext.compute(img3, kp2)
            t = Tracker(kp1 , desc1 , kp2 , desc2 , 0.7)
            res = t.match_flann()

        if res is not None:
            nm = "../Documents/TargetData/"+str(itr)+str("noTarThres0.7.txt")
            f = open(nm, 'w')
            for p in res:
                x, y = np.int32(kp2[p[1]].pt)
                toFile =  str(x)+","+str(y)+"\n"
                f.write(toFile)
                r = int(p.size)+7
                cv2.circle(img2, (x, y), r, (0, 255, 0) , thickness = -1)
                
                cv2.imshow('find_obj SURF flann', img3)
                cv2.imshow('find_obj SURF flann3', img2)
            f.close()
               # print clock() - tt 

        itr+=1
        ch = 0xFF & cv2.waitKey(1)
        if ch == 27:
            break
        
    cv2.destroyAllWindows()  


    

    
        